// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/parquet/thrift/ThriftTransport.h>
#include <melon/init/init.h>
#include <gtest/gtest.h>

using namespace kumo::pollux;
using namespace kumo::pollux::dwio::common;
using namespace kumo::pollux::parquet::thrift;

class ThriftTransportTest : public testing::Test {
 protected:
  void SetUp() override {
    input_.resize(bufferSize_);
    output_.resize(bufferSize_);
    for (size_t i = 0; i < input_.size(); ++i) {
      input_[i] = static_cast<uint8_t>(i);
    }
  }

  void prepareThriftStreamingTransport() {
    inputStream_ = std::make_shared<SeekableArrayInputStream>(
        input_.data(), input_.size(), 20);
    int32_t batchSize_;
    const void* bufferPointer;
    if (!inputStream_->Next(&bufferPointer, &batchSize_)) {
      POLLUX_CHECK(false, "Reading past end");
    }
    bufferStart_ = static_cast<const char*>(bufferPointer);
    bufferEnd_ = bufferStart_ + batchSize_;
    transport_ = std::make_shared<ThriftStreamingTransport>(
        inputStream_.get(), bufferStart_, bufferEnd_);
  }

  void prepareThriftBufferedTransport() {
    transport_ =
        std::make_shared<ThriftBufferedTransport>(input_.data(), bufferSize_);
  }

  static constexpr uint32_t bufferSize_ = 200;
  static constexpr uint32_t batchSize_ = 20;
  std::vector<uint8_t> input_;
  std::vector<uint8_t> output_;
  const char* bufferStart_{nullptr};
  const char* bufferEnd_{nullptr};
  std::shared_ptr<SeekableInputStream> inputStream_;
  std::shared_ptr<ThriftTransport> transport_;
};

TEST_F(ThriftTransportTest, streaming) {
  prepareThriftStreamingTransport();
  transport_->read(output_.data(), 10);
  transport_->read(output_.data() + 10, 50);
  transport_->read(output_.data() + 60, 140);

  for (size_t i = 0; i < input_.size(); ++i) {
    POLLUX_CHECK_EQ(input_[i], output_[i]);
  }
}

TEST_F(ThriftTransportTest, streamingOutOfBoundry) {
  prepareThriftStreamingTransport();
  transport_->read(output_.data(), 10);
  transport_->read(output_.data() + 10, 50);
  transport_->read(output_.data() + 60, 140);

  // The whole inputStream_ is consumed.
  EXPECT_ANY_THROW(transport_->read(output_.data() + bufferSize_, 1));
}

TEST_F(ThriftTransportTest, buffered) {
  prepareThriftBufferedTransport();
  transport_->read(output_.data(), 10);
  transport_->read(output_.data() + 10, 50);
  transport_->read(output_.data() + 60, 140);

  for (size_t i = 0; i < input_.size(); ++i) {
    POLLUX_CHECK_EQ(input_[i], output_[i]);
  }
}

TEST_F(ThriftTransportTest, bufferedOutOfBoundry) {
  prepareThriftStreamingTransport();
  transport_->read(output_.data(), 10);
  transport_->read(output_.data() + 10, 50);
  transport_->read(output_.data() + 60, 140);

  // The whole inputStream_ is consumed.
  EXPECT_ANY_THROW(transport_->read(output_.data() + bufferSize_, 1));
}

int main(int argc, char** argv) {
  testing::InitGoogleTest(&argc, argv);
  melon::Init init{&argc, &argv, false};
  return RUN_ALL_TESTS();
}
